"""
自动更新所有训练脚本以支持JSON日志
"""
import os
import re


def update_train_script(model_name):
    """
    更新单个模型的训练脚本

    Args:
        model_name: 模型名称
    """
    train_file = f'baselines/{model_name}/train.py'

    if not os.path.exists(train_file):
        print(f"[WARNING] {train_file} 不存在，跳过")
        return False

    print(f"[INFO] 更新 {train_file}...")

    with open(train_file, 'r', encoding='utf-8') as f:
        content = f.read()

    # 检查是否已经更新过
    if 'create_json_logger' in content and 'json_logger.log_epoch' in content:
        print(f"[INFO] {train_file} 已经包含JSON日志代码，跳过")
        return True

    # 1. 更新导入语句
    if 'create_json_logger' not in content:
        # 找到 from utils import 的位置
        import_pattern = r'from utils import \((.*?)\)'
        match = re.search(import_pattern, content, re.DOTALL)

        if match:
            imports = match.group(1)
            if 'create_json_logger' not in imports:
                # 在 setup_logger 后面添加
                new_imports = imports.replace(
                    'setup_logger, log_config,',
                    'setup_logger, log_config,\n    create_json_logger,'
                )
                content = content.replace(match.group(0), f'from utils import ({new_imports})')
                print(f"  ✓ 添加了 create_json_logger 导入")

    # 2. 在 train_fold 函数中添加JSON日志创建
    # 查找 logger = setup_logger 的位置
    logger_pattern = r"(logger = setup_logger\(log_dir, '[^']+', fold_idx\))"
    match = re.search(logger_pattern, content)

    if match and 'json_logger = create_json_logger' not in content:
        insert_pos = match.end()
        json_logger_code = f"\n\n    # 创建JSON日志记录器\n    json_logger = create_json_logger(log_dir, '{model_name}', fold_idx)"
        content = content[:insert_pos] + json_logger_code + content[insert_pos:]
        print(f"  ✓ 添加了JSON日志记录器创建代码")

    # 3. 在 log_config 后添加 json_logger.log_config
    log_config_pattern = r"(log_config\(logger, config_dict\))"
    match = re.search(log_config_pattern, content)

    if match and 'json_logger.log_config' not in content:
        insert_pos = match.end()
        json_config_code = "\n\n    # 记录配置到JSON日志\n    json_logger.log_config(config_dict)"
        content = content[:insert_pos] + json_config_code + content[insert_pos:]
        print(f"  ✓ 添加了JSON配置记录代码")

    # 4. 在训练循环开始添加 training_start_time
    train_start_pattern = r"(logger\.info\('开始训练\.\.\.'\))"
    match = re.search(train_start_pattern, content)

    if match and 'training_start_time = time.time()' not in content:
        insert_pos = match.end()
        start_time_code = "\n    training_start_time = time.time()"
        content = content[:insert_pos] + start_time_code + content[insert_pos:]
        print(f"  ✓ 添加了训练开始时间记录")

    # 5. 在训练循环中添加JSON日志记录
    # 这部分比较复杂，需要找到合适的位置插入
    # 简化处理：在文件末尾添加说明注释
    if 'json_logger.log_epoch' not in content:
        print(f"  ⚠ 需要手动添加 json_logger.log_epoch 调用")
        print(f"    请参考 train_template.py 或 UPDATE_TRAIN_SCRIPTS.md")

    # 保存更新后的文件
    backup_file = train_file + '.backup'
    with open(backup_file, 'w', encoding='utf-8') as f:
        f.write(content)

    with open(train_file, 'w', encoding='utf-8') as f:
        f.write(content)

    print(f"[INFO] {train_file} 更新完成（备份保存为 {backup_file}）")
    return True


def main():
    """主函数"""
    print("=" * 80)
    print("自动更新训练脚本以支持JSON日志")
    print("=" * 80)
    print()

    models = ['unet', 'deeplabv3plus', 'fcn', 'pspnet', 'unetpp', 'segnet']

    print("注意：由于训练脚本的复杂性，此脚本只能完成部分自动更新。")
    print("建议使用以下方法之一：")
    print()
    print("方法1（推荐）：使用 train_all.py 进行训练")
    print("  python train_all.py --models unet --fold 0")
    print()
    print("方法2：手动更新训练脚本")
    print("  参考 UPDATE_TRAIN_SCRIPTS.md 中的详细说明")
    print()
    print("方法3：直接使用更新后的 train_template.py")
    print("  所有模型都可以通过 train_all.py 使用这个模板")
    print()

    choice = input("是否继续尝试自动更新？(y/n): ").strip().lower()

    if choice != 'y':
        print("已取消")
        return

    print()
    print("开始更新...")
    print()

    success_count = 0
    for model in models:
        if update_train_script(model):
            success_count += 1
        print()

    print("=" * 80)
    print(f"更新完成！成功更新 {success_count}/{len(models)} 个模型")
    print("=" * 80)
    print()
    print("重要提示：")
    print("1. 自动更新可能不完整，请检查生成的代码")
    print("2. 备份文件保存为 *.backup，如有问题可以恢复")
    print("3. 建议使用 train_all.py 进行训练，它已完全支持JSON日志")
    print()
    print("验证方法：")
    print("  python baselines/unet/train.py --fold 0 --epochs 2")
    print("  ls logs/unet/  # 检查是否生成了 .json 文件")
    print()


if __name__ == '__main__':
    main()
